#!/usr/bin/env python3
import os
import yaml
import copy
import numpy as np
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from std_srvs.srv import Trigger
from geometry_msgs.msg import PoseWithCovarianceStamped
from enum import Enum, auto

class IMUMode(Enum):
    STATIC_CHECK = auto()
    CALIBRATION = auto()
    CORRECTION = auto()

imu_pose_covariance = [
    1e6, 0,    0,    0,    0,    0,    # x
    0,    1e6, 0,    0,    0,    0,    # y
    0,    0,    1e6,  0,    0,    0,    # z
    0,    0,    0,    3.0e-4,  0,    0,    # roll
    0,    0,    0,    0,    3.0e-4,  0,    # pitch    
    0,    0,    0,    0,    0,    3.0e-4   # yaw
]

class IMUCalibCorrector(Node):
    def __init__(self):
        super().__init__('imu_calib_corrector')

        # ==== 参数 ====
        self.declare_parameter('imu_topic_raw', '/imu/data_raw')
        self.declare_parameter('imu_topic_corrected', '/imu/data')
        self.declare_parameter('imu_pose_topic', '/pos_imu')
        self.declare_parameter('imu_bias_yaml_file', '/home/robot/auto-mower-v2/src/system/calibration_corrector/imu_cc/config/imu_bias.yaml')
        self.declare_parameter('sample_count', 500)
        self.declare_parameter('motion_threshold', 0.05)  # m/s²
        self.declare_parameter('motion_window', 50)       # 连续静止帧数

        self.imu_topic_raw = self.get_parameter('imu_topic_raw').get_parameter_value().string_value
        self.imu_topic_corrected = self.get_parameter('imu_topic_corrected').get_parameter_value().string_value
        self.imu_pose_topic = self.get_parameter('imu_pose_topic').get_parameter_value().string_value
        self.imu_bias_yaml_file = self.get_parameter('imu_bias_yaml_file').get_parameter_value().string_value
        self.sample_count = self.get_parameter('sample_count').get_parameter_value().integer_value
        self.motion_threshold = self.get_parameter('motion_threshold').get_parameter_value().double_value
        self.motion_window = self.get_parameter('motion_window').get_parameter_value().integer_value
     
        # ==== 数据缓存 ====
        self.samples = []
        self.bias_gyro = np.zeros(3)
        self.bias_acc = np.zeros(3)
        # self.imu_data = None

        # ==== 静止检测 ====
        self.motion_counter = 0
        self.prev_acc = None

        # ==== 订阅与发布 ====
        self.sub = self.create_subscription(Imu, self.imu_topic_raw, self.imu_callback, 100)
        self.imu_pub = self.create_publisher(Imu, self.imu_topic_corrected, 100)
        self.imu_pose_pub = self.create_publisher(PoseWithCovarianceStamped, self.imu_pose_topic, 100)

        # ==== Service ====
        self.create_service(Trigger, 'imu/start_calibration', self.start_calibration_cb)

        # ==== 模式管理 ====
        # mode 可取: IMUMode.STATIC_CHECK / CALIBRATION / CORRECTION
        if os.path.exists(self.imu_bias_yaml_file):
            self.mode = IMUMode.CORRECTION
            self.load_bias()
            self.get_logger().info(f"Bias file found. Starting in mode: {self.mode.name}")
        else:
            self.mode = IMUMode.STATIC_CHECK
            self.get_logger().info(f"Bias file not found. Starting in mode: {self.mode.name}")

    # ==== 回调 ====
    def imu_callback(self, msg:Imu):
        if self.mode == IMUMode.STATIC_CHECK:
            self._handle_static_check(msg)
        elif self.mode == IMUMode.CALIBRATION:
            self._handle_calibration(msg)
        elif self.mode == IMUMode.CORRECTION:
            self._handle_correction(msg)


    # ==== 静止检测 ====
    def _handle_static_check(self, msg:Imu):
        acc = np.array([msg.linear_acceleration.x,
                        msg.linear_acceleration.y,
                        msg.linear_acceleration.z])
        if self.prev_acc is not None:
            diff = np.linalg.norm(acc - self.prev_acc)
            if diff < self.motion_threshold:
                self.motion_counter += 1
            else:
                self.motion_counter = 0
        self.prev_acc = acc

        if self.motion_counter >= self.motion_window:
            self.get_logger().info("IMU is static. Starting calibration...")
            self.samples = []
            self.mode = IMUMode.CALIBRATION

    # ==== 标定模式 ====
    def _handle_calibration(self, msg):
        self.samples.append(copy.deepcopy(msg))
        if len(self.samples) >= self.sample_count:
            self.compute_bias()
            self.save_bias()
            self.samples = []
            self.mode = IMUMode.CORRECTION
            self.get_logger().info(f"Calibration finished. Switching to mode: {self.mode.name}")

    # ==== 修正模式 ====
    def _handle_correction(self, msg:Imu):
        # ==== 应用 bias 修正 ====
        current_time = self.get_clock().now().to_msg()
        corrected = Imu()
        corrected.header = msg.header
        corrected.header.stamp = current_time

        corrected.angular_velocity.x = msg.angular_velocity.x - self.bias_gyro[0]
        corrected.angular_velocity.y = msg.angular_velocity.y - self.bias_gyro[1]
        corrected.angular_velocity.z = msg.angular_velocity.z - self.bias_gyro[2]

        corrected.linear_acceleration.x = msg.linear_acceleration.x - self.bias_acc[0]
        corrected.linear_acceleration.y = msg.linear_acceleration.y - self.bias_acc[1]
        corrected.linear_acceleration.z = msg.linear_acceleration.z - self.bias_acc[2]

        corrected.orientation = msg.orientation
        corrected.orientation_covariance = msg.orientation_covariance
        corrected.angular_velocity_covariance = msg.angular_velocity_covariance
        corrected.linear_acceleration_covariance = msg.linear_acceleration_covariance

        # ==== 姿态 ====
        pose_cov = PoseWithCovarianceStamped()
        pose_cov.header = msg.header
        pose_cov.header.stamp = current_time
        pose_cov.pose.pose.position.x = 0.0
        pose_cov.pose.pose.position.y = 0.0
        pose_cov.pose.pose.position.z = 0.0
        pose_cov.pose.pose.orientation = msg.orientation
        pose_cov.pose.covariance[:] = imu_pose_covariance

        self.imu_pub.publish(corrected)
        self.imu_pose_pub.publish(pose_cov)

    # ==== Service触发标定 ====
    def start_calibration_cb(self, request, response):
        self.get_logger().info("Calibration service called.")
        self.samples = []
        self.mode = IMUMode.CALIBRATION
        self.motion_counter = 0
        self.prev_acc = None
        response.success = True
        response.message = "Calibration started."
        return response

    # ==== 标定计算 ====
    def compute_bias(self):
        gx = np.array([m.angular_velocity.x for m in self.samples])
        gy = np.array([m.angular_velocity.y for m in self.samples])
        gz = np.array([m.angular_velocity.z for m in self.samples])
        ax = np.array([m.linear_acceleration.x for m in self.samples])
        ay = np.array([m.linear_acceleration.y for m in self.samples])
        az = np.array([m.linear_acceleration.z for m in self.samples])

        self.bias_gyro = np.array([gx.mean(), gy.mean(), gz.mean()])
        self.bias_acc = np.array([ax.mean(), ay.mean(), az.mean() - 9.81])

        self.get_logger().info(f"Gyro bias: {self.bias_gyro}")
        self.get_logger().info(f"Accel bias: {self.bias_acc}")

    # ==== 保存 / 加载 bias ====
    def save_bias(self):
        data = {
            'bias_gyro': self.bias_gyro.tolist(),
            'bias_acc': self.bias_acc.tolist()
        }
        with open(self.imu_bias_yaml_file, 'w') as f:
            yaml.dump(data, f)
        self.get_logger().info(f"Bias saved to {self.imu_bias_yaml_file}")

    def load_bias(self):
        if os.path.exists(self.imu_bias_yaml_file):
            with open(self.imu_bias_yaml_file, 'r') as f:
                data = yaml.safe_load(f)
                if data is None:
                    self.get_logger().warn(f"Bias YAML file {self.imu_bias_yaml_file} is empty. Using zeros.")
                    data = {}
                self.bias_gyro = np.array(data.get('bias_gyro', [0, 0, 0]))
                self.bias_acc = np.array(data.get('bias_acc', [0, 0, 0]))
            self.get_logger().info(f"Loaded bias from {self.imu_bias_yaml_file}")
        else:
            self.get_logger().warn("Bias YAML not found. Using zeros.")

   
    

def main(args=None):
    rclpy.init(args=args)
    node = IMUCalibCorrector()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
